import argparse
from ultralytics import YOLO

def parse_opt():
    parser = argparse.ArgumentParser()
    parser.add_argument('--imgsz', type=int, default=640, help='Image size')
    parser.add_argument('--batch', type=int, default=8, help='Batch size')
    parser.add_argument('--epochs', type=int, default=300, help='Number of epochs')
    parser.add_argument('--data', type=str, default='data.yaml', help='Path to dataset.yaml')
    parser.add_argument('--weights', type=str, default='runs/detect/yolov10_continue/weights/best.pt', help='Path to previous best weights')  # 修改为之前保存的最佳权重文件
    parser.add_argument('--name', type=str, default='yolov10_continue', help='Name for the run')
    parser.add_argument('--lr0', type=float, default=0.001, help='Initial learning rate')  # 修改学习率参数
    return parser.parse_args()

def main(opt):
    model = YOLO(opt.weights)
    model.train(
        data=opt.data,
        epochs=opt.epochs,
        batch=opt.batch,
        imgsz=opt.imgsz,
        device='cuda',
        name=opt.name,
        patience=300,
        lr0=opt.lr0
    )

if __name__ == "__main__":
    opt = parse_opt()
    main(opt)
